#!/usr/bin/env python
# -*- coding=utf-8 -*-

import sys
reload(sys)
sys.setdefaultencoding('utf-8')

import os
from utils import TitleParser

def stat_genes(path):
    """
    gene_dict = {
        gene: [F1,F2..]
    }
    """
    gene_dict = {}
    files = os.listdir(path)
    for file in files:
        file_path = os.path.join(path,file)
        familyid = file.split('.')[0]
        with open(file_path,'r') as indata:
            title = indata.readline()
            tp = TitleParser(title)
            for line in indata:
                line_list = line.strip().split("\t")
                genes = tp.get_field(line_list,'GeneName')
                for gene in genes.split(','):
                    if gene not in gene_dict:
                        gene_dict[gene] = [familyid]
                    elif familyid not in gene_dict[gene]:
                        gene_dict[gene].append(familyid)
    return gene_dict

def output_gene_stat(gene_dict,ofile):
    """
    Gene Share_num Share_families
    """
    with open(ofile,'w') as odata:
        odata.write('Gene\tShare_num\tShare_families\n')
        for gene in gene_dict:
            share_num = len(gene_dict[gene])
            share_families = ",".join(gene_dict[gene])
            oline = "{gene}\t{share_num}\t{share_families}\n".format(**locals())
            odata.write(oline)
    print("Family Share information output to {ofile}".format(ofile=ofile))

if __name__ == "__main__":
    path = sys.argv[1]
    ofile = sys.argv[2]
    gene_dict = stat_genes(path)
    output_gene_stat(gene_dict,ofile)